import redis
from langchain.memory import ConversationBufferMemory
from langchain.schema import HumanMessage, AIMessage
from typing import Dict, List, Any
import json
from datetime import timedelta

class RedisMemory(ConversationBufferMemory):
    """Conversation memory with Redis storage."""
    
    redis_client: redis.Redis = None  # type: ignore
    key_prefix: str = "chat_history"
    ttl: int = 86400
    session_id: str = "default_session"
    
    def __init__(self, redis_url: str, session_id: str = "default_session",
                 key_prefix: str = "chat_history", ttl: int = 86400, **kwargs):
        """
        Initialize RedisMemory.
        
        Args:
            redis_url: Redis connection URL
            session_id: Unique session identifier
            key_prefix: Prefix for Redis keys
            ttl: Time-to-live in seconds for stored data
        """
        super().__init__(**kwargs)
        self.redis_client = redis.from_url(redis_url)
        self.key_prefix = key_prefix
        self.ttl = ttl
        self.session_id = session_id
    
    def _get_key(self) -> str:
        """Generate Redis key with prefix."""
        return f"{self.key_prefix}:{self.session_id}"
    
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save conversation context to Redis."""
        # Extract input string from inputs
        input_str = inputs.get("input", "") if isinstance(inputs, dict) else str(inputs)
        
        # Call parent method with proper inputs
        super().save_context({"input": input_str}, outputs)
        
        # Serialize and store messages
        key = self._get_key()
        messages = []
        for msg in self.chat_memory.messages:
            if isinstance(msg, HumanMessage):
                messages.append({"type": "human", "content": msg.content})
            elif isinstance(msg, AIMessage):
                messages.append({"type": "ai", "content": msg.content})
        
        # 使用整数秒数作为TTL
        self.redis_client.setex(key, self.ttl, json.dumps(messages))
    
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Load conversation history from Redis."""
        key = self._get_key()
        
        # Retrieve and deserialize messages
        stored_data = self.redis_client.get(key)
        if stored_data:
            messages_data = json.loads(stored_data)
            self.chat_memory.clear()
            for msg_data in messages_data:
                if msg_data["type"] == "human":
                    self.chat_memory.add_message(HumanMessage(content=msg_data["content"]))
                elif msg_data["type"] == "ai":
                    self.chat_memory.add_message(AIMessage(content=msg_data["content"]))
        
        return super().load_memory_variables(inputs)
    
    def store(self, key: str, value: str) -> None:
        """Store a value in Redis with the given key."""
        full_key = f"{self.key_prefix}:{key}"
        self.redis_client.setex(full_key, self.ttl, value)
    
    def retrieve(self, key: str) -> str:
        """Retrieve a value from Redis by key."""
        full_key = f"{self.key_prefix}:{key}"
        value = self.redis_client.get(full_key)
        return value.decode('utf-8') if value else ""